{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "%set_env PYTORCH_ENABLE_MPS_FALLBACK=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp models.kan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# KAN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Kolmogorov-Arnold Networks (KANs) are an alternative to Multi-Layer Perceptrons (MLPs). This model uses KANs similarly as our MLP model.\n",
    "\n",
    "**References**\n",
    "- [Ziming Liu, Yixuan Wang, Sachin Vaidya, Fabian Ruehle, James Halverson, Marin Soljačić, Thomas Y. Hou, Max Tegmark. \"KAN: Kolmogorov–Arnold Networks\"](https://arxiv.org/html/2404.19756v1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Figure 1. KAN compared to MLP.](imgs_models/kan.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "import logging\n",
    "import warnings\n",
    "from fastcore.test import test_eq\n",
    "from nbdev.showdoc import show_doc\n",
    "from neuralforecast.common._model_checks import check_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "from typing import Optional, Union\n",
    "\n",
    "import math\n",
    "\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from neuralforecast.losses.pytorch import MAE\n",
    "from neuralforecast.common._base_model import BaseModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class KANLinear(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    KANLinear \n",
    "    \"\"\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        in_features,\n",
    "        out_features,\n",
    "        grid_size=5,\n",
    "        spline_order=3,\n",
    "        scale_noise=0.1,\n",
    "        scale_base=1.0,\n",
    "        scale_spline=1.0,\n",
    "        enable_standalone_scale_spline=True,\n",
    "        base_activation=torch.nn.SiLU,\n",
    "        grid_eps=0.02,\n",
    "        grid_range=[-1, 1],\n",
    "    ):\n",
    "        super(KANLinear, self).__init__()\n",
    "        self.in_features = in_features\n",
    "        self.out_features = out_features\n",
    "        self.grid_size = grid_size\n",
    "        self.spline_order = spline_order\n",
    "\n",
    "        h = (grid_range[1] - grid_range[0]) / grid_size\n",
    "        grid = (\n",
    "            (\n",
    "                torch.arange(-spline_order, grid_size + spline_order + 1) * h\n",
    "                + grid_range[0]\n",
    "            )\n",
    "            .expand(in_features, -1)\n",
    "            .contiguous()\n",
    "        )\n",
    "        self.register_buffer(\"grid\", grid)\n",
    "\n",
    "        self.base_weight = torch.nn.Parameter(torch.Tensor(out_features, in_features))\n",
    "        self.spline_weight = torch.nn.Parameter(\n",
    "            torch.Tensor(out_features, in_features, grid_size + spline_order)\n",
    "        )\n",
    "        if enable_standalone_scale_spline:\n",
    "            self.spline_scaler = torch.nn.Parameter(\n",
    "                torch.Tensor(out_features, in_features)\n",
    "            )\n",
    "\n",
    "        self.scale_noise = scale_noise\n",
    "        self.scale_base = scale_base\n",
    "        self.scale_spline = scale_spline\n",
    "        self.enable_standalone_scale_spline = enable_standalone_scale_spline\n",
    "        self.base_activation = base_activation()\n",
    "        self.grid_eps = grid_eps\n",
    "\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        torch.nn.init.kaiming_uniform_(self.base_weight, a=math.sqrt(5) * self.scale_base)\n",
    "        with torch.no_grad():\n",
    "            noise = (\n",
    "                (\n",
    "                    torch.rand(self.grid_size + 1, self.in_features, self.out_features)\n",
    "                    - 1 / 2\n",
    "                )\n",
    "                * self.scale_noise\n",
    "                / self.grid_size\n",
    "            )\n",
    "            self.spline_weight.data.copy_(\n",
    "                (self.scale_spline if not self.enable_standalone_scale_spline else 1.0)\n",
    "                * self.curve2coeff(\n",
    "                    self.grid.T[self.spline_order : -self.spline_order],\n",
    "                    noise,\n",
    "                )\n",
    "            )\n",
    "            if self.enable_standalone_scale_spline:\n",
    "                # torch.nn.init.constant_(self.spline_scaler, self.scale_spline)\n",
    "                torch.nn.init.kaiming_uniform_(self.spline_scaler, a=math.sqrt(5) * self.scale_spline)\n",
    "\n",
    "    def b_splines(self, x: torch.Tensor):\n",
    "        \"\"\"\n",
    "        Compute the B-spline bases for the given input tensor.\n",
    "\n",
    "        Args:\n",
    "            x (torch.Tensor): Input tensor of shape (batch_size, in_features).\n",
    "\n",
    "        Returns:\n",
    "            torch.Tensor: B-spline bases tensor of shape (batch_size, in_features, grid_size + spline_order).\n",
    "        \"\"\"\n",
    "        assert x.dim() == 2 and x.size(1) == self.in_features\n",
    "\n",
    "        grid: torch.Tensor = (\n",
    "            self.grid\n",
    "        )  # (in_features, grid_size + 2 * spline_order + 1)\n",
    "        x = x.unsqueeze(-1)\n",
    "        bases = ((x >= grid[:, :-1]) & (x < grid[:, 1:])).to(x.dtype)\n",
    "        for k in range(1, self.spline_order + 1):\n",
    "            bases = (\n",
    "                (x - grid[:, : -(k + 1)])\n",
    "                / (grid[:, k:-1] - grid[:, : -(k + 1)])\n",
    "                * bases[:, :, :-1]\n",
    "            ) + (\n",
    "                (grid[:, k + 1 :] - x)\n",
    "                / (grid[:, k + 1 :] - grid[:, 1:(-k)])\n",
    "                * bases[:, :, 1:]\n",
    "            )\n",
    "\n",
    "        assert bases.size() == (\n",
    "            x.size(0),\n",
    "            self.in_features,\n",
    "            self.grid_size + self.spline_order,\n",
    "        )\n",
    "        return bases.contiguous()\n",
    "\n",
    "    def curve2coeff(self, x: torch.Tensor, y: torch.Tensor):\n",
    "        \"\"\"\n",
    "        Compute the coefficients of the curve that interpolates the given points.\n",
    "\n",
    "        Args:\n",
    "            x (torch.Tensor): Input tensor of shape (batch_size, in_features).\n",
    "            y (torch.Tensor): Output tensor of shape (batch_size, in_features, out_features).\n",
    "\n",
    "        Returns:\n",
    "            torch.Tensor: Coefficients tensor of shape (out_features, in_features, grid_size + spline_order).\n",
    "        \"\"\"\n",
    "        assert x.dim() == 2 and x.size(1) == self.in_features\n",
    "        assert y.size() == (x.size(0), self.in_features, self.out_features)\n",
    "\n",
    "        A = self.b_splines(x).transpose(\n",
    "            0, 1\n",
    "        )  # (in_features, batch_size, grid_size + spline_order)\n",
    "        B = y.transpose(0, 1)  # (in_features, batch_size, out_features)\n",
    "        solution = torch.linalg.lstsq(\n",
    "            A, B\n",
    "        ).solution  # (in_features, grid_size + spline_order, out_features)\n",
    "        result = solution.permute(\n",
    "            2, 0, 1\n",
    "        )  # (out_features, in_features, grid_size + spline_order)\n",
    "\n",
    "        assert result.size() == (\n",
    "            self.out_features,\n",
    "            self.in_features,\n",
    "            self.grid_size + self.spline_order,\n",
    "        )\n",
    "        return result.contiguous()\n",
    "\n",
    "    @property\n",
    "    def scaled_spline_weight(self):\n",
    "        return self.spline_weight * (\n",
    "            self.spline_scaler.unsqueeze(-1)\n",
    "            if self.enable_standalone_scale_spline\n",
    "            else 1.0\n",
    "        )\n",
    "\n",
    "    def forward(self, x: torch.Tensor):\n",
    "        assert x.dim() == 2 and x.size(1) == self.in_features\n",
    "\n",
    "        base_output = F.linear(self.base_activation(x), self.base_weight)\n",
    "        spline_output = F.linear(\n",
    "            self.b_splines(x).view(x.size(0), -1),\n",
    "            self.scaled_spline_weight.view(self.out_features, -1),\n",
    "        )\n",
    "        return base_output + spline_output\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def update_grid(self, x: torch.Tensor, margin=0.01):\n",
    "        assert x.dim() == 2 and x.size(1) == self.in_features\n",
    "        batch = x.size(0)\n",
    "\n",
    "        splines = self.b_splines(x)  # (batch, in, coeff)\n",
    "        splines = splines.permute(1, 0, 2)  # (in, batch, coeff)\n",
    "        orig_coeff = self.scaled_spline_weight  # (out, in, coeff)\n",
    "        orig_coeff = orig_coeff.permute(1, 2, 0)  # (in, coeff, out)\n",
    "        unreduced_spline_output = torch.bmm(splines, orig_coeff)  # (in, batch, out)\n",
    "        unreduced_spline_output = unreduced_spline_output.permute(\n",
    "            1, 0, 2\n",
    "        )  # (batch, in, out)\n",
    "\n",
    "        # sort each channel individually to collect data distribution\n",
    "        x_sorted = torch.sort(x, dim=0)[0]\n",
    "        grid_adaptive = x_sorted[\n",
    "            torch.linspace(\n",
    "                0, batch - 1, self.grid_size + 1, dtype=torch.int64, device=x.device\n",
    "            )\n",
    "        ]\n",
    "\n",
    "        uniform_step = (x_sorted[-1] - x_sorted[0] + 2 * margin) / self.grid_size\n",
    "        grid_uniform = (\n",
    "            torch.arange(\n",
    "                self.grid_size + 1, dtype=torch.float32, device=x.device\n",
    "            ).unsqueeze(1)\n",
    "            * uniform_step\n",
    "            + x_sorted[0]\n",
    "            - margin\n",
    "        )\n",
    "\n",
    "        grid = self.grid_eps * grid_uniform + (1 - self.grid_eps) * grid_adaptive\n",
    "        grid = torch.concatenate(\n",
    "            [\n",
    "                grid[:1]\n",
    "                - uniform_step\n",
    "                * torch.arange(self.spline_order, 0, -1, device=x.device).unsqueeze(1),\n",
    "                grid,\n",
    "                grid[-1:]\n",
    "                + uniform_step\n",
    "                * torch.arange(1, self.spline_order + 1, device=x.device).unsqueeze(1),\n",
    "            ],\n",
    "            dim=0,\n",
    "        )\n",
    "\n",
    "        self.grid.copy_(grid.T)\n",
    "        self.spline_weight.data.copy_(self.curve2coeff(x, unreduced_spline_output))\n",
    "\n",
    "    def regularization_loss(self, regularize_activation=1.0, regularize_entropy=1.0):\n",
    "        l1_fake = self.spline_weight.abs().mean(-1)\n",
    "        regularization_loss_activation = l1_fake.sum()\n",
    "        p = l1_fake / regularization_loss_activation\n",
    "        regularization_loss_entropy = -torch.sum(p * p.log())\n",
    "        return (\n",
    "            regularize_activation * regularization_loss_activation\n",
    "            + regularize_entropy * regularization_loss_entropy\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "\n",
    "class KAN(BaseModel):\n",
    "    \"\"\" KAN\n",
    "\n",
    "    Simple Kolmogorov-Arnold Network (KAN).\n",
    "    This network uses the Kolmogorov-Arnold approximation theorem, where splines\n",
    "    are learned to approximate more complex functions. Unlike the MLP, the\n",
    "    non-linear function are learned at the edges, and the nodes simply sum\n",
    "    the different learned functions.\n",
    "\n",
    "    **Parameters:**<br>\n",
    "    `h`: int, forecast horizon.<br>\n",
    "    `input_size`: int, considered autorregresive inputs (lags), y=[1,2,3,4] input_size=2 -> lags=[1,2].<br>\n",
    "    `grid_size`: int, number of intervals used by the splines to approximate the function.<br>\n",
    "    `spline_order`: int, order of the B-splines.<br>\n",
    "    `scale_noise`: float, regularization coefficient for the splines.<br>\n",
    "    `scale_base`: float, scaling coefficient for the base function.<br>\n",
    "    `scale_spline`: float, scaling coefficient for the splines.<br>\n",
    "    `enable_standalone_scale_spline`: bool, whether each spline is scaled individually.<br>\n",
    "    `grid_eps`: float, used for numerical stability.<br>\n",
    "    `grid_range`: list, range of the grid used for spline approximation.<br>\n",
    "    `n_hidden_layers`: int, number of hidden layers for the KAN.<br>\n",
    "    `hidden_size`: int or list, number of units for each hidden layer of the KAN. If an integer, all hidden layers will have the same size. Use a list to specify the size of each hidden layer.<br>    \n",
    "    `stat_exog_list`: str list, static exogenous columns.<br>\n",
    "    `hist_exog_list`: str list, historic exogenous columns.<br>\n",
    "    `futr_exog_list`: str list, future exogenous columns.<br>\n",
    "    `exclude_insample_y`: bool=False, the model skips the autoregressive features y[t-input_size:t] if True.<br>\n",
    "    `loss`: PyTorch module, instantiated train loss class from [losses collection](https://nixtla.github.io/neuralforecast/losses.pytorch.html).<br>\n",
    "    `valid_loss`: PyTorch module=`loss`, instantiated valid loss class from [losses collection](https://nixtla.github.io/neuralforecast/losses.pytorch.html).<br>\n",
    "    `max_steps`: int=1000, maximum number of training steps.<br>\n",
    "    `learning_rate`: float=1e-3, Learning rate between (0, 1).<br>\n",
    "    `num_lr_decays`: int=-1, Number of learning rate decays, evenly distributed across max_steps.<br>\n",
    "    `early_stop_patience_steps`: int=-1, Number of validation iterations before early stopping.<br>\n",
    "    `val_check_steps`: int=100, Number of training steps between every validation loss check.<br>\n",
    "    `batch_size`: int=32, number of different series in each batch.<br>\n",
    "    `valid_batch_size`: int=None, number of different series in each validation and test batch, if None uses batch_size.<br>\n",
    "    `windows_batch_size`: int=1024, number of windows to sample in each training batch, default uses all.<br>\n",
    "    `inference_windows_batch_size`: int=-1, number of windows to sample in each inference batch, -1 uses all.<br>\n",
    "    `start_padding_enabled`: bool=False, if True, the model will pad the time series with zeros at the beginning, by input size.<br>\n",
    "    `step_size`: int=1, step size between each window of temporal data.<br>\n",
    "    `scaler_type`: str='identity', type of scaler for temporal inputs normalization see [temporal scalers](https://nixtla.github.io/neuralforecast/common.scalers.html).<br>\n",
    "    `random_seed`: int=1, random_seed for pytorch initializer and numpy generators.<br>\n",
    "    `drop_last_loader`: bool=False, if True `TimeSeriesDataLoader` drops last non-full batch.<br>\n",
    "    `alias`: str, optional,  Custom name of the model.<br>\n",
    "    `optimizer`: Subclass of 'torch.optim.Optimizer', optional, user specified optimizer instead of the default choice (Adam).<br>\n",
    "    `optimizer_kwargs`: dict, optional, list of parameters used by the user specified `optimizer`.<br>\n",
    "    `dataloader_kwargs`: dict, optional, list of parameters passed into the PyTorch Lightning dataloader by the `TimeSeriesDataLoader`. <br>\n",
    "    `**trainer_kwargs`: int,  keyword trainer arguments inherited from [PyTorch Lighning's trainer](https://pytorch-lightning.readthedocs.io/en/stable/api/pytorch_lightning.trainer.trainer.Trainer.html?highlight=trainer).<br>  \n",
    "\n",
    "    **References**<br>\n",
    "    - [Ziming Liu, Yixuan Wang, Sachin Vaidya, Fabian Ruehle, James Halverson, Marin Soljačić, Thomas Y. Hou, Max Tegmark. \"KAN: Kolmogorov-Arnold Networks\"](https://arxiv.org/abs/2404.19756)\n",
    "    \"\"\"\n",
    "\n",
    "    # Class attributes\n",
    "    EXOGENOUS_FUTR = True\n",
    "    EXOGENOUS_HIST = True\n",
    "    EXOGENOUS_STAT = True    \n",
    "    MULTIVARIATE = False    # If the model produces multivariate forecasts (True) or univariate (False)\n",
    "    RECURRENT = False       # If the model produces forecasts recursively (True) or direct (False)\n",
    "\n",
    "    def __init__(self,\n",
    "                 h,\n",
    "                 input_size,\n",
    "                 grid_size: int = 5,\n",
    "                 spline_order: int = 3,\n",
    "                 scale_noise: float = 0.1,\n",
    "                 scale_base: float = 1.0,\n",
    "                 scale_spline: float = 1.0,\n",
    "                 enable_standalone_scale_spline: bool = True,\n",
    "                 grid_eps: float = 0.02,\n",
    "                 grid_range: list = [-1, 1],\n",
    "                 n_hidden_layers: int = 1,\n",
    "                 hidden_size: Union[int, list] = 512,\n",
    "                 stat_exog_list = None,\n",
    "                 hist_exog_list = None,\n",
    "                 futr_exog_list = None,\n",
    "                 exclude_insample_y = False,\n",
    "                 loss = MAE(),\n",
    "                 valid_loss = None,\n",
    "                 max_steps: int = 1000,\n",
    "                 learning_rate: float = 1e-3,\n",
    "                 num_lr_decays: int = -1,\n",
    "                 early_stop_patience_steps: int =-1,\n",
    "                 val_check_steps: int = 100,\n",
    "                 batch_size: int = 32,\n",
    "                 valid_batch_size: Optional[int] = None,\n",
    "                 windows_batch_size = 1024,\n",
    "                 inference_windows_batch_size = -1,\n",
    "                 start_padding_enabled = False,\n",
    "                 step_size: int = 1,\n",
    "                 scaler_type: str = 'identity',\n",
    "                 random_seed: int = 1,\n",
    "                 drop_last_loader: bool = False,\n",
    "                 alias: Optional[str] = None,\n",
    "                 optimizer = None,\n",
    "                 optimizer_kwargs = None,\n",
    "                 dataloader_kwargs = None,\n",
    "                 **trainer_kwargs):\n",
    "        \n",
    "        # Inherit BaseWindows class\n",
    "        super(KAN, self).__init__(h=h,\n",
    "                                  input_size=input_size,\n",
    "                                  futr_exog_list=futr_exog_list,\n",
    "                                  hist_exog_list=hist_exog_list,\n",
    "                                  stat_exog_list=stat_exog_list,\n",
    "                                  exclude_insample_y = exclude_insample_y,\n",
    "                                  loss=loss,\n",
    "                                  valid_loss=valid_loss,\n",
    "                                  max_steps=max_steps,\n",
    "                                  learning_rate=learning_rate,\n",
    "                                  num_lr_decays=num_lr_decays,\n",
    "                                  early_stop_patience_steps=early_stop_patience_steps,\n",
    "                                  val_check_steps=val_check_steps,\n",
    "                                  batch_size=batch_size,\n",
    "                                  valid_batch_size=valid_batch_size,\n",
    "                                  windows_batch_size=windows_batch_size,\n",
    "                                  inference_windows_batch_size=inference_windows_batch_size,\n",
    "                                  start_padding_enabled=start_padding_enabled,\n",
    "                                  step_size=step_size,\n",
    "                                  scaler_type=scaler_type,\n",
    "                                  drop_last_loader=drop_last_loader,\n",
    "                                  alias=alias,\n",
    "                                  random_seed=random_seed,\n",
    "                                  optimizer=optimizer,\n",
    "                                  optimizer_kwargs=optimizer_kwargs,\n",
    "                                  dataloader_kwargs = dataloader_kwargs,\n",
    "                                  **trainer_kwargs)\n",
    "        \n",
    "        # Architecture\n",
    "        self.n_hidden_layers = n_hidden_layers\n",
    "        self.hidden_size = hidden_size\n",
    "\n",
    "        input_size_first_layer = input_size + self.hist_exog_size * input_size + \\\n",
    "                                 self.futr_exog_size*(input_size + h) + self.stat_exog_size\n",
    "\n",
    "        if isinstance(self.hidden_size, int):\n",
    "            self.hidden_layers = [input_size_first_layer] + self.n_hidden_layers*[self.hidden_size] + [self.h * self.loss.outputsize_multiplier]\n",
    "        elif isinstance(self.hidden_size, list):\n",
    "            if len(self.hidden_size) != self.n_hidden_layers:\n",
    "                raise Exception(\"The number of elements in the list hidden_size must equal the number of n_hidden_layers\")\n",
    "            self.hidden_layers = [input_size_first_layer] + self.hidden_size + [self.h * self.loss.outputsize_multiplier]\n",
    "            \n",
    "        self.grid_size = grid_size\n",
    "        self.spline_order = spline_order\n",
    "        self.scale_noise = scale_noise\n",
    "        self.scale_base = scale_base\n",
    "        self.scale_spline = scale_spline\n",
    "        self.enable_standalone_scale_spline = enable_standalone_scale_spline\n",
    "        self.base_activation = torch.nn.SiLU\n",
    "        self.grid_eps = grid_eps\n",
    "        self.grid_range = grid_range\n",
    "\n",
    "        self.layers = torch.nn.ModuleList()\n",
    "        for in_features, out_features in zip(self.hidden_layers, self.hidden_layers[1:]):\n",
    "            self.layers.append(\n",
    "                KANLinear(\n",
    "                    in_features,\n",
    "                    out_features,\n",
    "                    grid_size=grid_size,\n",
    "                    spline_order=self.spline_order,\n",
    "                    scale_noise=self.scale_noise,\n",
    "                    scale_base=self.scale_base,\n",
    "                    scale_spline=self.scale_spline,\n",
    "                    base_activation=self.base_activation,\n",
    "                    grid_eps=self.grid_eps,\n",
    "                    grid_range=self.grid_range,\n",
    "                )\n",
    "            )\n",
    "\n",
    "\n",
    "    def regularization_loss(self, regularize_activation=1.0, regularize_entropy=1.0):\n",
    "        return sum(\n",
    "            layer.regularization_loss(regularize_activation, regularize_entropy)\n",
    "            for layer in self.layers\n",
    "        )\n",
    "        \n",
    "    def forward(self, windows_batch, update_grid=False):\n",
    "\n",
    "        insample_y = windows_batch['insample_y'].squeeze(-1)\n",
    "        futr_exog     = windows_batch['futr_exog']\n",
    "        hist_exog     = windows_batch['hist_exog']\n",
    "        stat_exog     = windows_batch['stat_exog']\n",
    "\n",
    "        # Flatten MLP inputs [B, L+H, C] -> [B, (L+H)*C]\n",
    "        # Contatenate [ Y_t, | X_{t-L},..., X_{t} | F_{t-L},..., F_{t+H} | S ]\n",
    "        batch_size = len(insample_y)\n",
    "        if self.hist_exog_size > 0:\n",
    "            insample_y = torch.cat(( insample_y, hist_exog.reshape(batch_size,-1) ), dim=1)\n",
    "\n",
    "        if self.futr_exog_size > 0:\n",
    "            insample_y = torch.cat(( insample_y, futr_exog.reshape(batch_size,-1) ), dim=1)\n",
    "\n",
    "        if self.stat_exog_size > 0:\n",
    "            insample_y = torch.cat(( insample_y, stat_exog.reshape(batch_size,-1) ), dim=1)\n",
    "\n",
    "        y_pred = insample_y.clone()\n",
    "        for layer in self.layers:\n",
    "            if update_grid:\n",
    "                layer.update_grid(y_pred)\n",
    "            y_pred = layer(y_pred)\n",
    "\n",
    "        y_pred = y_pred.reshape(batch_size, self.h, \n",
    "                                self.loss.outputsize_multiplier)\n",
    "        return y_pred\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(KAN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(KAN.fit, name='KAN.fit')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(KAN.predict, name='KAN.predict')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Unit tests for models\n",
    "logging.getLogger(\"pytorch_lightning\").setLevel(logging.ERROR)\n",
    "logging.getLogger(\"lightning_fabric\").setLevel(logging.ERROR)\n",
    "with warnings.catch_warnings():\n",
    "    warnings.simplefilter(\"ignore\")\n",
    "    check_model(KAN, checks=[\"airpassengers\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Usage Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| eval: false\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from neuralforecast import NeuralForecast\n",
    "from neuralforecast.models import KAN\n",
    "from neuralforecast.losses.pytorch import DistributionLoss\n",
    "from neuralforecast.utils import AirPassengersPanel, AirPassengersStatic\n",
    "\n",
    "Y_train_df = AirPassengersPanel[AirPassengersPanel.ds<AirPassengersPanel['ds'].values[-12]] # 132 train\n",
    "Y_test_df = AirPassengersPanel[AirPassengersPanel.ds>=AirPassengersPanel['ds'].values[-12]].reset_index(drop=True) # 12 test\n",
    "\n",
    "fcst = NeuralForecast(\n",
    "    models=[\n",
    "            KAN(h=12,\n",
    "                input_size=24,\n",
    "                loss = DistributionLoss(distribution=\"Normal\"),\n",
    "                max_steps=100,\n",
    "                scaler_type='standard',\n",
    "                futr_exog_list=['y_[lag12]'],\n",
    "                hist_exog_list=None,\n",
    "                stat_exog_list=['airline1'],\n",
    "                ),     \n",
    "    ],\n",
    "    freq='ME'\n",
    ")\n",
    "fcst.fit(df=Y_train_df, static_df=AirPassengersStatic)\n",
    "forecasts = fcst.predict(futr_df=Y_test_df)\n",
    "\n",
    "# Plot quantile predictions\n",
    "Y_hat_df = forecasts.reset_index(drop=False).drop(columns=['unique_id','ds'])\n",
    "plot_df = pd.concat([Y_test_df, Y_hat_df], axis=1)\n",
    "plot_df = pd.concat([Y_train_df, plot_df])\n",
    "\n",
    "plot_df = plot_df[plot_df.unique_id=='Airline1'].drop('unique_id', axis=1)\n",
    "plt.plot(plot_df['ds'], plot_df['y'], c='black', label='True')\n",
    "plt.plot(plot_df['ds'], plot_df['KAN-median'], c='blue', label='median')\n",
    "plt.fill_between(x=plot_df['ds'][-12:], \n",
    "                 y1=plot_df['KAN-lo-90'][-12:].values,\n",
    "                 y2=plot_df['KAN-hi-90'][-12:].values,\n",
    "                 alpha=0.4, label='level 90')\n",
    "plt.legend()\n",
    "plt.grid()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
